Speed up CI: fix NY tax-benefit-system cloning, run batches concurrently, rebalance to 14 runners#8886
Conversation
|
Round 2 pushed — re-batched from the measured per-batch peak RSS that round 1's new logging produced. Round 1 result: 9/15 heavy jobs passed with large wins (rest 33.7→10.4 min, Partners 29.8→18.0, irs-household 23.6→13.9), but 2-wide batch co-scheduling OOM'd four jobs: the RSS telemetry revealed single batches already peak at 12-14.6 GB on Linux ( Round 2 is headroom-first, per maintainer direction (long-term robustness to new tests > minimum runner count):
Now 23 full-suite runners (was 17 before this PR), each with ≥ ~8 GB of slack for future heavy tests, and every log permanently reports per-batch peak RSS + the 25 slowest cases so drift is visible before it breaks. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #8886 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 3 3
Lines 55 99 +44
Branches 0 10 +10
=========================================
+ Hits 55 99 +44
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Data point for this rebalance from tonight's merge queue: "Full Suite - Contrib (states-shard-2)" was killed (exit 143 / runner-canceled, ~17-26 min in) on at least nine separate PR runs across #8679, #8818, #8901 (twice), #8902, #8903 (twice), #8907, and #8908 — every one passed unchanged on rerun, so it's load/memory, not tests. That shard looks like the first candidate for whatever split this PR lands on. |
The NY EITC and pre-2024 CTC formulas deep-cloned the entire tax-benefit system (parameter tree + variable registry) on every call that hit the decoupled/pre-TCJA branch, driving the NY credits test memory to ~12 GB and requiring a dedicated quarantined CI runner. Build each pinned system (pre-ARPA 2020 EITC, pre-TCJA 2017 CTC) once per process in a module-level cache keyed on the baseline system's identity, so layered reforms still rebuild against their own base. Measured on the credits test files: ny_ctc_pre_2024.yaml 373s -> 57s, ny_eitc.yaml 106s -> 53s, peak RSS 4.8 -> 2.8 GB, outputs identical. Fixes PolicyEngine#8114 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
New --workers N flag (default 1 = unchanged sequential behavior) runs batch subprocesses concurrently, longest-first by YAML file count, with per-batch output buffered and printed atomically plus a 60s heartbeat. Each batch now reports its subprocess peak RSS (VmHWM from /proc on Linux) so CI logs show real per-batch memory. Grace sleep after pytest completion reduced from 5s to 1s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With the NY clone fix, states-ny folds back into the states matrix (8 batches, 2 shards). Reform and hhs merge into the baseline-contrib runner, usda moves to the ssa runner, contrib other-shards 1+3 merge, and the partners job fans analytics_coverage/edge_cases out per topic/state folder (invocation-only; partner files untouched). Heavy targets run batches 2-wide, the light rest job 3-wide. All test jobs set PYTEST_ADDOPTS to print the 25 slowest cases and skip pytest's unraisableexception gc sweep (~11s per subprocess). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
test_microsim.py now builds one Microsimulation per dataset via a module-scoped fixture shared across the parametrized years (2 full builds instead of 4; ~17 min of the Rest CI job). The RUN_HEAVY_TESTS LSR/CG interaction tests (still skipped by default) subsample to 10,000 households and reuse a single baseline, so opting in is cheaper. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Round-1's --workers 2 co-scheduling OOM-killed four runners: the new per-batch RSS logging showed single batches already peak far higher on Linux than local estimates suggested (ri/ctc_reform_test.yaml 14.6 GB alone, contrib/states/oh 12.4 GB, baseline/contrib/states 12.2 GB, gov/simulation 12.1 GB, half of baseline/household 11.8 GB). Headroom-first re-layout: every heavy target runs one batch at a time (--workers 2 kept only for partners, max batch 2.5 GB), oversized batches are split finer (states --batches 16 over 4 shards, household --batches 4, gov/simulation and baseline/contrib/states per-file, oh joins ri in PER_FILE_STATES), the two files too big for any batching are split by reform combo (RI CTC into 3, CRFB payroll into 2, cases byte-identical), and the microsimulation pytest run gets its own job since the combined process was OOM-killed at 94%. Every subprocess now stays at or below ~8 GB, leaving half the 16 GB runner free for future test growth. Quick Feedback's selective cap drops to 100 files (one subprocess with 250 yaml files exceeds runner memory; over-cap safely defers to the full suites). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Quick Feedback OOM'd twice: run_selective_tests packed every selected location into one process (this PR selected the whole NY folder plus all RI and CRFB reform files), where reformed tax-benefit systems and per-case simulations accumulate for the process lifetime — exceeding runner memory under coverage even though each location alone is fine. One subprocess per location bounds memory at the heaviest single location; coverage -a appends results across invocations unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MaxGhenis
left a comment
MaxGhenis
left a comment
There was a problem hiding this comment.
Fable review + instrumented verification: the pinned-TBS cache replaces per-call full-system clones in the NY pre-TCJA CTC and pre-ARPA EITC formulas (12→2 clones, peak shard RSS 7.30→3.72 GB, byte-identical outputs), with weakref identity keying so reformed systems rebuild rather than leak baseline parameters. Shard rebalance proven exhaustive: 4,280/4,280 test files covered, zero orphans/dups. This addresses the exit-143 kills that hit ~20 runs in the last day. CI green.
Folder-per-batch with a hand-maintained heavy-state list OOM'd within a day of PR #8886: two merged PRs (#8816, #8856) added reform test files to 13 state folders, and the GA batch silently grew from 3 to 11 distinct reform combos. Each distinct (reforms, parameter-overrides) combination deep-copies the full tax-benefit system into policyengine-core's never-evicted cache (~1.45 GB each, calibrated against 46 contrib batches on CI run 28756084246), so that one subprocess peaked at 15.2 GB on 16 GB runners and killed contrib states-shard-4 on most runs after 2026-07-05 23:00 UTC. - test_batched.py now packs each state's files into batches capped at reform-combo weight 5 (~8.7 GB predicted peak), replacing the PER_FILE_STATES hand-list and the RI shard-pinning special case — both subsumed by the general mechanism. - Split the four over-budget files (OR/DE at 7 combos, OH/MD at 6) into two files each by combo group; cases relocated byte-identically. - New code-health test caps the combo weight one test file may carry, so future over-budget files fail at authoring time instead of OOM-ing CI later. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2 participants
Summary
CI took ~34 min wall-clock across 17 runners, sat at the memory cliff (single batches measured at 12-14.6 GB on 16 GB runners — the source of the historical intermittent "runner lost communication" failures), and needed hand-rearranging as tests grew. This PR fixes the worst per-case cost bug, re-batches the suite from measured per-batch peak RSS, and adds permanent timing/memory telemetry to every CI log.
Final state (validated by 3 full runs on this PR): all green, zero OOMs, ~26 min wall on 23 full-suite runners, every subprocess ≤ ~8.5 GB peak — half of each runner's memory left as headroom so newly added heavy tests don't require re-tuning.
Fixes #8114
Changes
1. NY #8114 fix (
policyengine_us/tools/pinned_tbs.py+ 3 NY credit variables)The NY EITC / pre-2024 CTC formulas deep-cloned the entire tax-benefit system per formula call. Each pinned system (pre-ARPA 2020 EITC, pre-TCJA 2017 CTC) is now built once per process in a module-level cache keyed on the baseline system's identity (weakref — layered reforms rebuild against their own base). Measured:
ny_ctc_pre_2024.yaml373s → 57s,ny_eitc.yaml106s → 53s, peak RSS 4.8 → 2.8 GB, outputs identical. The dedicatedstates-nyCI runner is retired; NY runs in the ordinary states shards.2. Permanent CI telemetry
PYTEST_ADDOPTS="--durations=25 -p no:unraisableexception -p no:threadexception"— each log ends with the 25 slowest test cases, and subprocesses skip pytest 8.4's ~11s full-heap gc sweep at exit.test_batched.pyreports each batch subprocess's true peak RSS (VmHWMfrom/proc) in its footer and the summary. Heavy new tests become visible in green runs, not red ones.test_batched.pyalso gained an opt-in--workers Npool; CI uses 2 only for the Partners job, whose batches max at 2.5 GB. Everything else runs single-worker after measurement showed reform-heavy batches co-scheduled 2-wide exceed 16 GB.)3. Memory-headroom re-batching (measured, not estimated)
Per-batch peak RSS from this PR's own runs drove the layout: states
--batches 16across 4 shards, household--batches 4,gov/simulationandbaseline/contrib/statesper-file,ohjoinsriin contrib per-file isolation, congress/irs/ssa heavy batches run strictly solo. Two files were too big for any batching and are split by reform-parameter combination with byte-identical cases:contrib/states/ri/ctc_reform_test.yaml(14.6 GB alone) → 3 files;contrib/crfb/tax_employer_payroll_tax_percentage.yaml(9.6 GB) → 2 files.4. Process-lifetime memory bugs in test runners
make test-otherpytest process OOM'd at 94%: microsimulation tests now run as their own CI job, andtest_microsim.pybuilds one Microsimulation per dataset (module-scoped fixture; 2 builds instead of 4 — it was the real ~17-min hotspot of the old Rest job). TheRUN_HEAVY_TESTSLSR/CG tests (skipped on CI) now subsample and share a baseline.run_selective_tests.py(Quick Feedback) packed every selected location into one process — reuniting the very files split above and OOMing under coverage. It now runs one location per subprocess (coverage run -aappends across invocations), bounding memory at the heaviest single location.5. Job layout: 17 → 23 runners
Baseline: states-shard-1..4, irs, household, ssa-usda, rest-a, rest-b, contrib-hhs, reform. Contrib: states-shard-1..4, other-shard-1, 2a, 2b, 3, congress. Plus Partners, Rest (Python + variables), Microsimulation. Total runner-minutes are roughly unchanged — the compute is spread wider and shorter, with growth headroom on every runner.
Measured results (run on commit 1d4e2b9)
All 35 checks green. Wall 26m27s (was ~33.7 min). Old→new highlights: Baseline rest 33.7 → 12.0∥10.6; Rest+microsim 29.8 → 11.8∥11.7; irs-household 23.6 → 14.5∥10.9; Partners 29.8 → 22.2; contrib other max 26.4 → 14.1; reform 9.6 → 9.4. Worst job: contrib states-shard-4 (26m27s — carries the RI per-file batches).
Notes for maintainers
Baseline (states-ny/states-non-ny-*/rest/contrib/ssa/irs-household),Reform (per-file),Contrib (other-shard-1/2/3 as previously named); new names per the layout above.Test plan
🤖 Generated with Claude Code